#include <stdlib.h>
#include <vector>

#include <cppunit/extensions/HelperMacros.h>

#include "TestPlatform.h"
#include "../common/TestUtils.h"

#include "../../FormatString.h"
#include "../../LapackPlatform/Lapack.h"
#include "../../Profiler.h"


using namespace std;

struct Matrix {
    int ld;
    float *data;	

    Matrix(float *values, int dim) : data(values), ld(dim) { }

    inline float& operator () (int row, int col) const {
        return data[row + col * ld];
    }

    inline float* column(int col) const {
        return &data[col * ld];
    }

    inline Matrix submatrix(int row, int col) const {
        return Matrix(&data[col * ld + row], ld);
    }
};


inline void CHECK_LAPACK(int info) {
	CPPUNIT_ASSERT(info == 0);
}

void multiplySquare(const Matrix& result, int size, const Matrix& m1, const Matrix& m2) {
    for (int row = 0; row < size; row++) {
        for (int col = 0; col < size; col++) {
            float s = 0;
            for (int i = 0; i < size; i++) {
                s += m1(row, i) * m2(i, col);
            }
            result(row, col) = s;
        }
    }
}

void transpose(const Matrix& m) {
    for (int row = 0; row < m.ld; row++) {
        for (int col = row + 1; col < m.ld; col++) {
            std::swap(m(row, col), m(col, row));
        }
    }
}

void fillTridiagonal(Matrix& m, int size, float *diagonal, float *offDiagonal) {
    for (int i = 0; i < size - 1; i++) {
        m(i, i) = diagonal[i];
        m(i, i + 1) = offDiagonal[i];
        m(i + 1, i) = offDiagonal[i];
    }
    m(size - 1, size - 1) = diagonal[size - 1];
}

TestPlatform::TestPlatform(Lapack& _lapack) : lapack(_lapack) { }

void TestPlatform::testTridiagonalization(integer size) {

	vector<real> matrix(size * size);
	Matrix m(&matrix[0], size);
	for (integer i = 0; i < size; i++) {
		for (integer j = i; j < size; j++) {
			m(i, j) = m(j, i) = getRandom();
		}
	}

	vector<real> src(matrix);

	vector<real> d(size);
	vector<real> e(size);
	vector<real> tau(size);

	integer minusOne = -1;
	integer lwork = -1;
	integer info = 0;

	real temp;    
	lapack.ssytrd("L", &size, NULL, &size, NULL, NULL, NULL, &temp, &minusOne, &info);
	lwork = (integer)temp;

	lapack.sorgtr("L", &size, NULL, &size, NULL, &temp, &minusOne, &info);
	lwork = max(lwork, (integer)temp);

	vector<real> work(lwork);    

	lapack.ssytrd("L", &size, &matrix[0], &size, &d[0], &e[0], &tau[0], &work[0], &lwork, &info);	
	CHECK_LAPACK(info);
	
	lapack.sorgtr("L", &size, &matrix[0], &size, &tau[0], &work[0], &lwork, &info);    
	CHECK_LAPACK(info);
	
	vector<real> vTri(size * size);
	Matrix tridiagonal(&vTri[0], size);

	fillTridiagonal(tridiagonal, size, &d[0], &e[0]);

	Matrix transform(&matrix[0], size);
	vector<real> vRes(size * size);
	Matrix res(&vRes[0], size);

	multiplySquare(res, size, transform, tridiagonal);
	transpose(transform);
	multiplySquare(tridiagonal, size, res, transform);

	CPPUNIT_ASSERT( compare(&src[0], tridiagonal.data, size * size, 1E-4f) == 0);
}

